#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct  3 14:12:53 2019

@author: yl
"""
import json
from PIL import Image
import torch
from torch import nn
from torchvision import models,datasets,transforms
from torch.autograd import Variable
import os
from efficientnet_pytorch import EfficientNet
import matplotlib.pyplot as plt
import torch.nn.functional as F


transform=transforms.Compose([
            transforms.Resize([224,224]),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],
                                 std=[0.229,0.224,0.225])
                            ])


def prediect(img_path):
    '''
    model_name='resnet50'
    model=EfficientNet.from_name(model_name)
    feature=model._fc.in_features
    model._fc=nn.Linear(in_features=feature,out_features=10,bias=True)
    '''
    
    model=models.resnet101(pretrained=True)
    for parma in model.parameters():
        parma.requires_grad=False
    model.fc=torch.nn.Linear(2048,10)   
    resume='model_best.path.tar'
    checkpoint=torch.load(resume)
    model.load_state_dict(checkpoint['state_dict'])
    model=model.cuda()


    with torch.no_grad():
        model.eval()
    img=Image.open(img_path)
    img=transform(img).unsqueeze(0)
    img_ = Variable(img.cuda())
    outputs = model(img_)
    _, indices = torch.max(outputs, 1)
    percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100

    labels_map = json.load(open('labels_fish_identification.txt'))
    labels_map = [labels_map[str(i)] for i in range(10)]

    perc = percentage[int(indices)].item()
    result = labels_map[indices]

    print("Prediction : ", result, ", Score: ", percentage)

    _, indices = torch.sort(outputs, descending=True)  # obtain predicted result and sort form big to small
    percentage1 = torch.nn.functional.softmax(outputs, dim=1)[0] * 100  # reture percentage of predicted value
    # print([(labels_map[idx], percentage1[idx].item()) for idx in indices[0][:1]])
    result = [(labels_map[idx], percentage1[idx].item()) for idx in indices[0][:1]]
    print(result[0])
    print(result[0][0])
    print(result[0][1])

if __name__ == '__main__':
    prediect('2.jpg')
    
####  https://blog.csdn.net/heiheiya/article/details/103031300  code of this link can display result in image, such probability and class 